package com.flame.auth.client.mybatis;

import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import com.flame.auth.core.annotation.DataFilter;
import com.flame.auth.core.annotation.DataFilters;
import com.flame.auth.core.constants.AuthConstants;
import com.flame.auth.core.context.AuthContextHolder;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

/**
 * DataAuthInterceptor
 *
 * @author wuxintong
 * @since 2023/2/11
 */
@Intercepts({@Signature(type = Executor.class, method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query",
                args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class,
                        BoundSql.class})})
public class DataAuthInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];

        if (!SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType())) {
            return invocation.proceed();
        }
        // 1、获取 DataFilter 注解信息
        List<DataFilter> dataFilterList = getDataFilterList(mappedStatement);
        if (CollectionUtils.isEmpty(dataFilterList)) {
            return invocation.proceed();
        }

        // 2、准备拼接条件
        Map<String, String> dataAuthMap = AuthContextHolder.getContext().getDataAuthMap();
        List<String> conditionList = new ArrayList<>();
        for (DataFilter dataFilter : dataFilterList) {
            if (AuthConstants.DATA_TYPE_USER.equals(dataFilter.dataType())) {
                String userDataAuth = dataAuthMap.get(AuthConstants.DATA_TYPE_USER);
                if (StringUtils.hasText(userDataAuth)) {
                    if (AuthConstants.DATA_TYPE_USER_ALL.equals(userDataAuth)) {
                        conditionList.add("1 = 1");
                    } else {
                        conditionList.add(dataFilter.column() + userDataAuth);
                    }
                }
            }
            if (AuthConstants.DATA_TYPE_ORG.equals(dataFilter.dataType())) {
                String orgDataAuth = dataAuthMap.get(AuthConstants.DATA_TYPE_ORG);
                if (StringUtils.hasText(orgDataAuth)) {
                    conditionList.add(dataFilter.column() + orgDataAuth);
                } else {
                    return new ArrayList<>();
                }
            }
        }
        String conditionStr = String.join(" and ", conditionList);
        if (!StringUtils.hasText(conditionStr)) {
            return new ArrayList<>();
        }

        // 3、开始拼sql
        BoundSql boundSql = mappedStatement.getBoundSql(invocation.getArgs()[1]);
        if (invocation.getArgs().length == 6) {
            boundSql = (BoundSql) invocation.getArgs()[5];
        }
        String oldSql = boundSql.getSql();
        Select select = (Select) CCJSqlParserUtil.parse(oldSql);
        PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
        Expression where = plainSelect.getWhere();
        if (where != null) {
            Expression left = CCJSqlParserUtil.parseCondExpression("(" + where + ")");
            Expression right = CCJSqlParserUtil.parseCondExpression("(" + conditionStr + ")");
            plainSelect.setWhere(new AndExpression(left, right));
        } else {
            plainSelect.setWhere(CCJSqlParserUtil.parseCondExpression(conditionStr));
        }
        String newSql = plainSelect.toString();

        // 3、设置新 sql
        MappedStatement newMappedStatement = newMappedStatement(mappedStatement, boundSql, newSql);
        invocation.getArgs()[0] = newMappedStatement;

        return invocation.proceed();
    }

    private MappedStatement newMappedStatement(MappedStatement mappedStatement, BoundSql boundSql, String newSql) {
        BoundSql newBoundSql = new BoundSql(mappedStatement.getConfiguration(), newSql, boundSql.getParameterMappings(),
                boundSql.getParameterObject());

        for (ParameterMapping parameterMapping : boundSql.getParameterMappings()) {
            String property = parameterMapping.getProperty();
            if (boundSql.hasAdditionalParameter(property)) {
                newBoundSql.setAdditionalParameter(property, boundSql.getAdditionalParameter(property));
            }
        }

        MappedStatement.Builder builder =
                new MappedStatement.Builder(mappedStatement.getConfiguration(), mappedStatement.getId(),
                        parameterObject -> newBoundSql, mappedStatement.getSqlCommandType());
        builder.resource(mappedStatement.getResource());
        builder.fetchSize(mappedStatement.getFetchSize());
        builder.statementType(mappedStatement.getStatementType());
        builder.keyGenerator(mappedStatement.getKeyGenerator());
        if (mappedStatement.getKeyProperties() != null && mappedStatement.getKeyProperties().length > 0) {
            builder.keyProperty(mappedStatement.getKeyProperties()[0]);
        }
        builder.timeout(mappedStatement.getTimeout());
        builder.parameterMap(mappedStatement.getParameterMap());
        builder.resultMaps(mappedStatement.getResultMaps());
        builder.resultSetType(mappedStatement.getResultSetType());
        builder.cache(mappedStatement.getCache());
        builder.flushCacheRequired(mappedStatement.isFlushCacheRequired());
        builder.useCache(mappedStatement.isUseCache());

        return builder.build();
    }

    private List<DataFilter> getDataFilterList(MappedStatement mappedStatement) {
        List<DataFilter> dataFilterList = new ArrayList<>();
        try {
            String id = mappedStatement.getId();
            String className = id.substring(0, id.lastIndexOf("."));
            String methodName = id.substring(id.lastIndexOf(".") + 1);
            Method[] methods = Class.forName(className).getMethods();
            for (Method method : methods) {
                if (method.getName().equals(methodName)) {
                    for (Annotation annotation : method.getDeclaredAnnotations()) {
                        if (annotation instanceof DataFilters) {
                            DataFilters dataFilters = (DataFilters) annotation;
                            dataFilterList.addAll(Arrays.asList(dataFilters.value()));
                        }
                        if (annotation instanceof DataFilter) {
                            dataFilterList.add((DataFilter) annotation);
                        }
                    }
                    break;
                }
            }
            return dataFilterList;
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        }
        return dataFilterList;
    }
}
